import tensorflow as tf
from data import *
from model import *
import matplotlib.pyplot as plt

# 读取数据集
train_dataset = read_Voc()
val_dataset = read_Voc(file_type='val')

BATCH_SIZE = 32
EPOCHS = 10
BUFFER_SIZE = len(train_dataset)
STEPS_PER_EPOCH = BUFFER_SIZE // BATCH_SIZE

VAL_SUBSPLITS = 5
VALIDATION_STEPS = len(val_dataset)//BATCH_SIZE//VAL_SUBSPLITS

# 打乱训练数据集
train_dataset = train_dataset.shuffle(buffer_size=BUFFER_SIZE)
train_dataset = train_dataset.repeat()
train_dataset = train_dataset.batch(BATCH_SIZE)
train_dataset = train_dataset.prefetch(buffer_size=AUTOTUNE)

val_dataset = val_dataset.batch(BATCH_SIZE)

model = UNetCompiled(input_size=(128, 128, 3), n_filters=32, n_classes=3)
model.compile(optimizer=tf.keras.optimizers.Adam(),
             loss='categorical_crossentropy',
             metrics=['accuracy'])

model_history = model.fit(
    train_dataset,
    epochs=EPOCHS,
    steps_per_epoch=STEPS_PER_EPOCH,
    validation_steps=VALIDATION_STEPS,
    validation_data=val_dataset,
)

loss = model_history.history['loss']
val_loss = model_history.history['val_loss']

epochs = range(EPOCHS)

plt.figure()
plt.plot(epochs, loss, 'r', label='Training loss')
plt.plot(epochs, val_loss, 'bo', label='Validation loss')
plt.title('Training and Validation Loss')
plt.xlabel('Epoch')
plt.ylabel('Loss Value')
plt.ylim([0, 1])
plt.legend()
plt.show()